import {
  createConnection,
  ProposedFeatures,
  TextDocuments,
  Diagnostic,
  DiagnosticSeverity,
  TextDocumentSyncKind,
  CompletionItem,
  CompletionItemKind,
  Hover,
  MarkupKind,
  Position,
  Range,
  Location,
  SymbolInformation,
  SymbolKind,
  CodeAction,
  CodeActionKind,
  TextEdit,
  SemanticTokensBuilder,
  SemanticTokenTypes,
  SemanticTokenModifiers,
} from "vscode-languageserver/node";
import { TextDocument } from "vscode-languageserver-textdocument";

// Create LSP connection over stdio
const connection = createConnection(ProposedFeatures.all);
// Manage open text documents
const documents: TextDocuments<TextDocument> = new TextDocuments(TextDocument);

connection.onInitialize(() => {
  return {
    capabilities: {
      textDocumentSync: TextDocumentSyncKind.Incremental,
      definitionProvider: true,
      referencesProvider: true,
      workspaceSymbolProvider: true,
      completionProvider: { resolveProvider: true, triggerCharacters: ["T", "F", "N"] },
      hoverProvider: true,
      codeActionProvider: true,
      documentFormattingProvider: true,
      semanticTokensProvider: {
        legend: {
          tokenTypes: [
            SemanticTokenTypes.keyword,
            SemanticTokenTypes.comment,
          ],
          tokenModifiers: [
            SemanticTokenModifiers.deprecated,
            SemanticTokenModifiers.modification,
            SemanticTokenModifiers.documentation,
            SemanticTokenModifiers.readonly,
          ],
        },
        full: true,
        range: true,
      },
    },
  };
});

// Validate on open and on change
documents.onDidOpen((e) => {
  validateTextDocument(e.document);
});
documents.onDidChangeContent((change) => {
  validateTextDocument(change.document);
});

// Optional: clear diagnostics on close
documents.onDidClose((e) => {
  connection.sendDiagnostics({ uri: e.document.uri, diagnostics: [] });
});

async function validateTextDocument(textDocument: TextDocument): Promise<void> {
  const text = textDocument.getText();
  const diagnostics: Diagnostic[] = [];

  // Rule: highlight all occurrences of "TODO" as Warning
  const todoRegex: RegExp = /TODO/g;
  let match: RegExpExecArray | null;
  while ((match = todoRegex.exec(text))) {
    const start = textDocument.positionAt(match.index);
    const end = textDocument.positionAt(match.index + match[0].length);
    diagnostics.push({
      range: { start, end },
      severity: DiagnosticSeverity.Warning,
      message: "Found TODO",
      source: "learn-lsp",
    });
  }

  const fixmeRegex: RegExp = /FIXME/g;
  let fixmeMatch: RegExpExecArray | null;
  while ((fixmeMatch = fixmeRegex.exec(text))) {
    const start = textDocument.positionAt(fixmeMatch.index);
    const end = textDocument.positionAt(
      fixmeMatch.index + fixmeMatch[0].length
    );
    diagnostics.push({
      range: { start, end },
      severity: DiagnosticSeverity.Error,
      message: "Found FIXME",
      source: "learn-lsp",
    });
  }

  const lines = text.split("\n");
  let charCount = 0;
  for (const line of lines) {
    if (line.length > 80) {
      const start = textDocument.positionAt(charCount);
      const end = textDocument.positionAt(charCount + line.length);
      diagnostics.push({
        range: { start, end },
        severity: DiagnosticSeverity.Warning,
        message: "LINE TOO LONG: " + String(line.length),
        source: "learn-lsp",
      });
    }
    charCount += line.length + 1; // include the newline when advancing
  }

  connection.sendDiagnostics({ uri: textDocument.uri, diagnostics });
}

// ---- Completion & Hover helpers ----
function getWordRange(doc: TextDocument, position: Position): { word: string; range: Range } | null {
  const text = doc.getText();
  const offset = doc.offsetAt(position);
  let start = offset;
  let end = offset;
  const isWordChar = (ch: string) => /[A-Za-z_]/.test(ch);
  while (start > 0 && isWordChar(text.charAt(start - 1))) start--;
  while (end < text.length && isWordChar(text.charAt(end))) end++;
  if (start === end) return null;
  const word = text.slice(start, end);
  return { word, range: { start: doc.positionAt(start), end: doc.positionAt(end) } };
}

function escapeRegExp(s: string): string {
  return s.replace(/[.*+?^${}()|[\]\\]/g, "\\$&");
}

function findWordOccurrences(doc: TextDocument, word: string): Range[] {
  const text = doc.getText();
  const re = new RegExp(`\\b${escapeRegExp(word)}\\b`, "g");
  const ranges: Range[] = [];
  let m: RegExpExecArray | null;
  while ((m = re.exec(text))) {
    const start = doc.positionAt(m.index);
    const end = doc.positionAt(m.index + m[0].length);
    ranges.push({ start, end });
  }
  return ranges;
}

connection.onCompletion((params) => {
  const doc = documents.get(params.textDocument.uri);
  if (!doc) return [];
  const wr = getWordRange(doc, params.position);
  const prefix = wr ? wr.word.toUpperCase() : "";
  const items: CompletionItem[] = [
    { label: "TODO", kind: CompletionItemKind.Keyword, detail: "Task marker" },
    { label: "FIXME", kind: CompletionItemKind.Keyword, detail: "Bug marker" },
    { label: "NOTE", kind: CompletionItemKind.Keyword, detail: "Note marker" },
  ];
  const filtered = prefix ? items.filter((i) => i.label.toUpperCase().startsWith(prefix)) : items;
  return filtered;
});

connection.onCompletionResolve((item) => {
  item.documentation = { kind: MarkupKind.Markdown, value: `**${item.label}** marker.` };
  return item;
});

connection.onHover((params) => {
  const doc = documents.get(params.textDocument.uri);
  if (!doc) return null;
  const wr = getWordRange(doc, params.position);
  if (!wr) return null;
  const W = wr.word.toUpperCase();
  if (W === "TODO") {
    const contents: Hover["contents"] = { kind: MarkupKind.Markdown, value: "Marks a task to be done." };
    return { contents, range: wr.range };
  }
  if (W === "FIXME") {
    const contents: Hover["contents"] = { kind: MarkupKind.Markdown, value: "Marks a known problem to be fixed." };
    return { contents, range: wr.range };
  }
  return null;
});

connection.onDefinition((params) => {
  const doc = documents.get(params.textDocument.uri);
  if (!doc) return null;
  const wr = getWordRange(doc, params.position);
  if (!wr) return null;
  const occ = findWordOccurrences(doc, wr.word);
  if (occ.length === 0) return null;
  const range = occ[0];
  const loc: Location = { uri: doc.uri, range };
  return loc;
});

connection.onReferences((params) => {
  const doc = documents.get(params.textDocument.uri);
  if (!doc) return [];
  const wr = getWordRange(doc, params.position);
  if (!wr) return [];
  const occ = findWordOccurrences(doc, wr.word);
  return occ.map((range) => ({ uri: doc.uri, range }));
});

connection.onWorkspaceSymbol((params) => {
  const query = (params.query || "").trim();
  const out: SymbolInformation[] = [];
  const docs = documents.all();
  for (const doc of docs) {
    if (query) {
      const ranges = findWordOccurrences(doc, query);
      for (const r of ranges) {
        out.push({ name: query, kind: SymbolKind.Constant, location: { uri: doc.uri, range: r } });
      }
    } else {
      for (const kw of ["TODO", "FIXME", "NOTE"]) {
        const ranges = findWordOccurrences(doc, kw);
        for (const r of ranges) {
          out.push({ name: kw, kind: SymbolKind.Constant, location: { uri: doc.uri, range: r } });
        }
      }
    }
  }
  return out;
});

// Start listening
documents.listen(connection);
connection.listen();

connection.onDocumentFormatting((params) => {
  const doc = documents.get(params.textDocument.uri);
  if (!doc) return [];
  const original = doc.getText();

  // Simple formatter: trim trailing spaces per line and ensure final newline
  const lines = original.split(/\r?\n/);
  const trimmed = lines.map((l) => l.replace(/[ \t]+$/g, ""));
  let formatted = trimmed.join("\n");
  if (!formatted.endsWith("\n")) {
    formatted = formatted + "\n";
  }

  if (formatted === original) {
    return [];
  }
  const full: TextEdit = {
    range: { start: doc.positionAt(0), end: doc.positionAt(original.length) },
    newText: formatted,
  };
  return [full];
});

connection.onCodeAction((params) => {
  const doc = documents.get(params.textDocument.uri);
  if (!doc) return [];

  const only = params.context.only || [];
  const quickfixRequested = only.length === 0 || only.includes(CodeActionKind.QuickFix);
  if (!quickfixRequested) return [];

  const actions: CodeAction[] = [];
  const uri = doc.uri;

  // Prefer using diagnostics from the client context if available
  const diags = params.context.diagnostics || [];
  for (const d of diags) {
    const textInRange = doc.getText(d.range).toUpperCase();
    if (textInRange === "TODO") {
      actions.push({
        title: "Mark TODO as DONE",
        kind: CodeActionKind.QuickFix,
        isPreferred: true,
        diagnostics: [d],
        edit: { changes: { [uri]: [{ range: d.range, newText: "DONE" }] } },
      });
      actions.push({
        title: "Remove TODO",
        kind: CodeActionKind.QuickFix,
        diagnostics: [d],
        edit: { changes: { [uri]: [{ range: d.range, newText: "" }] } },
      });
    }
  }

  // If no diagnostics were provided, attempt a word-based quick fix at the range start
  if (actions.length === 0) {
    const wr = getWordRange(doc, params.range.start);
    if (wr && wr.word.toUpperCase() === "TODO") {
      actions.push({
        title: "Mark TODO as DONE",
        kind: CodeActionKind.QuickFix,
        edit: { changes: { [uri]: [{ range: wr.range, newText: "DONE" }] } },
      });
      actions.push({
        title: "Remove TODO",
        kind: CodeActionKind.QuickFix,
        edit: { changes: { [uri]: [{ range: wr.range, newText: "" }] } },
      });
    }
  }

  return actions;
});


function positionLE(a: Position, b: Position): boolean {
  return a.line < b.line || (a.line === b.line && a.character <= b.character);
}

const TOKEN_TYPE_INDEX: Record<string, number> = {
  keyword: 0,
  comment: 1,
};
const MOD_INDEX: Record<string, number> = {
  deprecated: 0,
  modification: 1,
  documentation: 2,
  readonly: 3,
};
function toModifierBits(mods: string[]): number {
  let bits = 0;
  for (const m of mods) {
    const idx = MOD_INDEX[m];
    if (typeof idx === 'number') {
      bits |= (1 << idx);
    }
  }
  return bits;
}

function buildSemanticTokensForDoc(doc: TextDocument, range?: Range) {
  const text = doc.getText();
  const builder = new SemanticTokensBuilder();

  const pushWord = (word: string, tokenTypeName: keyof typeof TOKEN_TYPE_INDEX, mods: (keyof typeof MOD_INDEX)[]) => {
    const re = new RegExp(`\\b${escapeRegExp(word)}\\b`, "g");
    let m: RegExpExecArray | null;
    while ((m = re.exec(text))) {
      const startOffset = m.index;
      const endOffset = m.index + m[0].length;
      const startPos = doc.positionAt(startOffset);
      const endPos = doc.positionAt(endOffset);
      if (range) {
        if (!(positionLE(range.start, startPos) && positionLE(endPos, range.end))) {
          continue;
        }
      }
      const typeIndex = TOKEN_TYPE_INDEX[tokenTypeName as string];
      const modBits = toModifierBits(mods as string[]);
      builder.push(startPos.line, startPos.character, m[0].length, typeIndex, modBits);
    }
  };

  // Keywords and markers
  pushWord("TODO", "keyword", ["modification"]);
  pushWord("FIXME", "keyword", ["deprecated"]);
  pushWord("DONE", "keyword", ["readonly"]);
  pushWord("NOTE", "comment", ["documentation"]);

  return builder.build();
}

connection.languages.semanticTokens.on((params) => {
  const doc = documents.get(params.textDocument.uri);
  if (!doc) return { data: [] };
  return buildSemanticTokensForDoc(doc);
});

connection.languages.semanticTokens.onRange((params) => {
  const doc = documents.get(params.textDocument.uri);
  if (!doc) return { data: [] };
  return buildSemanticTokensForDoc(doc, params.range);
});